# !/usr/bin/env python
# -*-coding:utf-8 -*-

"""
# File       : infer.py
# Time       ：2021/12/20 16:05
# Author     ：caoxu
# version    ：python 3.9
# Description：
"""

import cv2
import os
import shutil
import numpy as np
import pandas as pd
import tensorflow as tf
import y3.utils as utils
from y3.config import cfg
from y3.yolov3 import YOLOv3, decode


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in ['.png', '.jpg', '.jpeg', '.PNG', '.JPG', '.JPEG'])


if __name__ == '__main__':
    # Build Model
    print('-----Infer Start-----')
    print('-----Load Config-----')
    INPUT_SIZE = cfg.TEST.INPUT_SIZE
    NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
    CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)

    print('-----Build Model-----')
    input_layer = tf.keras.layers.Input([INPUT_SIZE, INPUT_SIZE, 3])
    feature_maps = YOLOv3(input_layer, NUM_CLASS)
    bbox_tensors = []
    for i, fm in enumerate(feature_maps):
        bbox_tensor = decode(fm, i, NUM_CLASS)
        bbox_tensors.append(bbox_tensor)

    model = tf.keras.Model(input_layer, bbox_tensors)
    model.load_weights(cfg.TRAIN.WEIGHTS_DIR + 'yolov3.weights')       #   + 'Y3'

    print('-----Load Images-----')
    img_path = './test_images/111'
    image_filenames = [os.path.join(img_path, x) for x in os.listdir(img_path) if is_image_file(x)]
    for image_filename in image_filenames:
        print(image_filename)
        image = cv2.imread(image_filename)

        # Predict Process
        print('-----Predict Process-----')
        image_size = image.shape[:2]
        image_data = utils.image_preporcess(np.copy(image), [INPUT_SIZE, INPUT_SIZE])
        image_data = image_data[np.newaxis, ...].astype(np.float32)

        print('-----Predict bboxes NMS-----')
        pred_bbox = model.predict(image_data)
        pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
        pred_bbox = tf.concat(pred_bbox, axis=0)
        bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)
        bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')

        print('-----Draw and Save Result Image-----')
        image = utils.draw_bbox(image, bboxes)
        cv2.imwrite(image_filename + '_result.jpg', image)